package com.digiwin.demo.sink;

import com.digiwin.demo.constant.ConfigConstant;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.streaming.connectors.elasticsearch.ElasticsearchSinkFunction;
import org.apache.flink.streaming.connectors.elasticsearch.RequestIndexer;
import org.apache.flink.streaming.connectors.elasticsearch.util.RetryRejectedExecutionFailureHandler;
import org.apache.flink.streaming.connectors.elasticsearch6.ElasticsearchSink;
import org.apache.flink.streaming.connectors.elasticsearch6.RestClientFactory;
import org.apache.http.Header;
import org.apache.http.HttpHost;
import org.apache.http.message.BasicHeader;
import org.elasticsearch.client.RestClientBuilder;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.List;

/**
 * @Author yanggld
 * @Date 2019/12/26-15:58
 * @Description
 */
public abstract class AbstractElasticsearchSinkFunction<T> implements Serializable {

	private static final String IP_ADDRESS = ConfigConstant.ES_HOST;
	private static final int port = ConfigConstant.ES_PORT;
	private static final String scheme = ConfigConstant.ES_SCHEME;

	private ElasticsearchSink<T> elasticsearchSink;

	public AbstractElasticsearchSinkFunction() {
		List<HttpHost> httpHosts = new ArrayList<>();
		httpHosts.add(new HttpHost(IP_ADDRESS, port, scheme));
		// use a ElasticsearchSink.Builder to create an ElasticsearchSink
		ElasticsearchSink.Builder<T> esSinkBuilder = new ElasticsearchSink.Builder<>(httpHosts, (ElasticsearchSinkFunction<T>) (o, runtimeContext, requestIndexer) -> invoke(o,runtimeContext,requestIndexer));
		esSinkBuilder.setBulkFlushMaxActions(1);
		esSinkBuilder.setRestClientFactory(new RestClientFactoryImpl());
		esSinkBuilder.setFailureHandler(new RetryRejectedExecutionFailureHandler());
		this.elasticsearchSink = esSinkBuilder.build();
	}

	public ElasticsearchSink<T> getElasticsearchSink() {
		return elasticsearchSink;
	}

	public abstract void invoke(T value, RuntimeContext runtimeContext, RequestIndexer requestIndexer);

}

class RestClientFactoryImpl implements RestClientFactory {
	@Override
	public void configureRestClientBuilder(RestClientBuilder restClientBuilder) {
		Header[] headers = new BasicHeader[]{new BasicHeader("Content-Type","application/json")};
		restClientBuilder.setDefaultHeaders(headers); //以数组的形式可以添加多个header
		restClientBuilder.setMaxRetryTimeoutMillis(90000);
	}
}